from neutron import Plugin, User
from fastapi import WebSocket
from starlette import websockets
import asyncio
from typing import Literal
import json

app = Plugin(__file__)

LOG_MAX_LEN = 2000  # 保留日志长度
LOG_REMOVE_LEN = 2010  # 进行删除的长度

id_lst = 0  # 任务 id
ws_lst = 0  # 连接 id

tasks_list = {}
tasksobj_list = {}

websockets_lst_callbacks = {}
buttons_lst_callbacks = {}
events_lst = {}


def remove_attr(dct: dict, key: str):
    dct_ret = dct.copy()
    del dct_ret[key]
    return dct_ret


class TasksTools:
    def __init__(self, set_id, user: User) -> None:
        self.id = set_id
        self.user = user

    async def set_process(self, process: int):  # 更改进度
        global tasks_list, tasksobj_list, id_lst
        if (tasks_list[self.id]["status"] != "running"):  # 如果当前没在运行就无法改变
            return
        tasks_list[self.id]["progress"] = process  # 更新进度
        for i in websockets_lst_callbacks.values():  # 更新所有连接的状态
            await (i("update", {
                "id": self.id,
                "data": remove_attr(tasks_list[self.id], "user")
            }, self.user))

    async def write_log(self, log: str):  # 写入日志
        global tasks_list, tasksobj_list, id_lst
        if (tasks_list[self.id]["status"] != "running"):
            return
        tasks_list[self.id]["log"] += log
        # 删除日志
        if (len(tasks_list[self.id]["log"]) > LOG_REMOVE_LEN):
            tasks_list[self.id]["log"] = tasks_list[self.id]["log"][len(
                tasks_list[self.id]["log"])-LOG_MAX_LEN:]
        for i in websockets_lst_callbacks.values():
            await i("update", {
                "id": self.id,
                "data": remove_attr(tasks_list[self.id], "user")
            }, self.user)

    async def exit(self, exit_mode: bool = False):  # 正常退出
        global tasks_list, tasksobj_list, id_lst
        if (tasks_list[self.id]["status"] != "running"):
            return
        if (exit_mode):  # 是否正常退出
            tasks_list[self.id]["status"] = "failed"
        else:
            tasks_list[self.id]["status"] = "success"
        for i in websockets_lst_callbacks.values():
            await i("update", {
                "id": self.id,
                "data": remove_attr(tasks_list[self.id], "user")
            }, self.user)
        # 停止当前任务
        await stop_task(self.id, self.user, tasks_list[self.id]["status"])

    # 设置按钮列表
    async def show_buttons(self, button_lists: list[str]):
        global tasks_list, tasksobj_list, id_lst,  buttons_lst_callbacks
        events_lst[self.id] = asyncio.Event()
        return_val = -1
        if (tasks_list[self.id]["status"] != "running"):
            return
        tasks_list[self.id]["buttons"] = button_lists
        for i in websockets_lst_callbacks.values():
            await (i("update", {
                "id": self.id,
                "data": remove_attr(tasks_list[self.id], "user")
            }, self.user))

        async def buttons_callback(button_id: int, ev_id: int, user: User, jump_check_user=False):
            nonlocal return_val
            if (not jump_check_user and not await app.checker.check_func_permission(user, 2, "admin") and self.user != user):
                return
            global buttons_lst_callbacks
            return_val = button_id
            events_lst[ev_id].set()

        buttons_lst_callbacks[self.id] = buttons_callback

        await events_lst[self.id].wait()
        del events_lst[self.id]
        tasks_list[self.id]["buttons"] = []
        del buttons_lst_callbacks[self.id]
        for i in websockets_lst_callbacks.values():
            await (i("update", {
                "id": self.id,
                "data": remove_attr(tasks_list[self.id], "user")
            }, self.user))

        return return_val


async def ignore(*args, **kwargs):
    pass


@app.export
async def create_task(name: str, description: str, command, user: User):  # 创建任务
    global id_lst, tasks_list
    tasks_list[id_lst] = {
        "name": name,
        "description": description,
        "log": "",
        "status": "running",
        "progress": 0,
        "buttons": [],
        "user": user
    }

    tasksobj_list[id_lst] = asyncio.create_task(
        command(TasksTools(id_lst, user)))  # 启动任务

    for i in websockets_lst_callbacks.values():  # 更新连接
        await (i("add", {
            "id": id_lst,
            "data": remove_attr(tasks_list[id_lst], "user")
        }, user))
    id_lst += 1
    return id_lst-1


async def stop_task(task_id: int, user: User, status: str = "stop"):  # 停止任务
    if (task_id in tasks_list):
        if (tasks_list[task_id]["user"] != user and not await app.checker.check_func_permission(user, 2, "admin")):
            return
        if (task_id in buttons_lst_callbacks):
            del buttons_lst_callbacks[task_id]
        tasks_list[task_id]["buttons"] = []
        tasksobj_list[task_id].cancel()  # 停止 task
        tasks_list[task_id]["status"] = status
        for i in websockets_lst_callbacks.values():
            await (i("update", {
                "id": task_id,
                "data": remove_attr(tasks_list[task_id], "user")
            }, user))


async def remove_task(task_id: int, user: User):  # 删除任务
    if (task_id in tasks_list):
        if (tasks_list[task_id]["user"] != user and not await app.checker.check_func_permission(user, 2, "admin")):
            return
        if (tasks_list[task_id]["status"] == "running"):
            await stop_task(task_id, user)
        for i in websockets_lst_callbacks.values():
            await (i("remove", {
                "id": task_id
            }, user))
        del tasksobj_list[task_id]
        del tasks_list[task_id]


@app.route.ws("/tasks", level=1)
async def tasks_ws(user: User, websocket: WebSocket, show_all=False):  # 主连接
    global ws_lst
    if (show_all == True and not await app.checker.check_func_permission(user, 2, "admin")):
        show_all = False
    now_ws = ws_lst
    ws_lst += 1
    websockets_lst_callbacks[now_ws] = ignore
    try:
        await websocket.accept()
        await websocket.send_json({  # 初始化客户端
            "type": "init",
            "data": {k: remove_attr(v, "user") for k, v in tasks_list.items() if v["user"] == user or show_all}
        })

        # 设置回调
        async def write_ws(type: Literal['add', 'remove', 'update'], data: object, user_get: User):
            if (user_get != user):
                return
            await websocket.send_text(json.dumps({
                "type": type,
                "data": data
            }))
        websockets_lst_callbacks[now_ws] = write_ws

        while True:
            data = await websocket.receive_json()
            if (data["type"] == "refresh"):
                await websocket.send_text(json.dumps({
                    "type": "init",
                    "data": {k: remove_attr(v, "user") for k, v in tasks_list.items() if v["user"] == user or show_all}
                }))
            elif (data["type"] == "stop"):
                await (stop_task(data["data"]["id"], user))
            elif (data["type"] == "remove"):
                await (remove_task(data["data"]["id"], user))
            elif (data["type"] == "button"):
                if (data["data"]["id"] in buttons_lst_callbacks):
                    await (
                        buttons_lst_callbacks[data["data"]["id"]](data["data"]["data"], data["data"]["id"], user))
    except websockets.WebSocketDisconnect:
        pass
    finally:
        del websockets_lst_callbacks[now_ws]  # 删除回调


# @app.route.post("/test")
# async def test():  # 测试用 route
#     async def command(tools: TasksTools):
#         print("Return:", await tools.show_buttons(["你好", "测试"]))
#         for i in range(1, 51):
#             await tools.set_process(i*2)
#             await tools.write_log(f"Load: {i}\n")
#             await asyncio.sleep(0.2)
#         await tools.write_log(f"Finish\n")
#         await tools.exit()
#     await create_task("你好", "test", command)
